import os
import random
import math
os.environ['KMP_DUPLICATE_LIB_OK']='True'
import torch
import torch.nn as nn
class FixedDownsampleConv(nn.Module):
    def __init__(self, input_height, input_width, output_height=224, output_width=224,
                 in_channels=3, out_channels=3, kernel_size=3, padding=0):
        super().__init__()
        self.kernel_size = kernel_size
        self.padding = padding

        # 严格根据卷积公式计算步长，确保输出尺寸正确
        self.stride_h = input_height // output_height
        self.stride_w = input_width // output_width

        # 处理无法整除的情况（微调步长确保输出尺寸正确）
        # 若整除有误差，可适当增加步长1
        if (input_height % output_height) > 0:
            self.stride_h += 1
        if (input_width % output_width) > 0:
            self.stride_w += 1

        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=(self.stride_h, self.stride_w),
            padding=padding,
            groups=in_channels
        )

    def calculate_stride(self, input_size, output_size):
        """根据输入/输出尺寸和卷积参数，计算合适的步长"""
        # 公式变形：stride ≤ (input_size - kernel_size + 2×padding) / (output_size - 1)
        max_stride = (input_size - self.kernel_size + 2 * self.padding) / (output_size - 1)
        # 步长必须为正整数，取不超过max_stride的最大整数
        stride = int(max_stride)
        # 确保步长至少为1
        return max(stride, 1)

    def forward(self, x):
        out = self.conv(x)
        return out


# 测试：输入4000×2400，目标输出224×224
if __name__ == "__main__":
    input_h, input_w = 896, 896
    model = FixedDownsampleConv(input_h, input_w)

    input_tensor = torch.randn(2, 3, input_h, input_w)
    output_tensor = model(input_tensor)

    print(f"输入尺寸: {input_h}×{input_w}")
    print(f"使用步长: ({model.stride_h}, {model.stride_w})")
    print(f"输出尺寸: {output_tensor.shape[2]}×{output_tensor.shape[3]}")  # 应为224×224